import sensor, image, time, uos
from machine import Pin

# 摄像头初始化
sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QVGA)  # 320x240
sensor.skip_frames(time = 2000)    # 让摄像头稳定

# SD卡保存路径
save_dir = "/sd/dataset"
try:
    uos.mkdir(save_dir)
except:
    pass

# 按键初始化（假设按键接在PIO16上）
key = Pin(16, Pin.IN, Pin.PULL_UP)

def key_pressed():
    return key.value() == 0  # 按下返回True

# 拍照计数
img_count = 0

while(True):
    try:
        img = sensor.snapshot()
        img.draw_string(10, 10, "Press KEY to Capture", color=(255,0,0))

        if key_pressed():
            # 清理内存
            import gc
            gc.collect()
            filename = "{}/img_{}.jpg".format(save_dir, img_count)
            img.save(filename)
            print("[SUCCESS] Saved: {}".format(filename))
            img_count += 1

            # 显示保存成功
            img.draw_string(10, 30, "Saved: {}".format(filename), color=(0,255,0))
            time.sleep_ms(500)  # 防抖延时

    except Exception as e:
        print("[ERROR]", e)
        time.sleep(1)
